#include "../Utils.cuh"

using namespace OBR;

/**
 * @brief calc combined effect of bsdf / cos theta
 *
 */
static __forceinline__ __device__ Spectrum calcFr(SamplePRD *prd, float *pdf,
                                                  MixedMaterialData *data)
{
    *pdf = 0.0f;
    Spectrum fr = Spectrum(0.0f);
    float prev_weight = 0.0f;
    for (size_t i = 0; i < data->num_mat_mixed; i++)
    {
        float m_pdf;
        float weight = data->weights[i] - prev_weight;
        prd->inter.mat_instance_id = data->mat_instance_ids[i];
        fr += optixDirectCall<Spectrum, SamplePRD *, float *, float, float>(
                  data->mat_ids[i] + prd->params->material_base_offset, prd, &m_pdf, 0, 0,
                  MaterialCallMode::CALC_FR_PDF) *
              weight;
        *pdf += m_pdf * weight;
    }
    return fr * data->scale;
}

/**
 * @brief generate new ray info and add to the next ray, then return beta of the ray
 *
 * @param prd
 * @return beta
 */
static __forceinline__ __device__ Spectrum rayGeneration(SamplePRD *prd, float *pdf,
                                                         MixedMaterialData *data, float u0,
                                                         float u1)
{
    // uniform sample one type of material
    float prev = 0.0f;
    uint mat = 0;
    for (size_t i = 0; i < data->num_mat_mixed; i++)
    {
        mat = i;
        if (data->weights[i] > u0)
            break;
        prev = data->weights[i];
    }
    float p = data->weights[mat] - prev;
    float u = (u0 - prev) / p;

    prd->inter.mat_instance_id = data->mat_instance_ids[mat];
    Spectrum fr = optixDirectCall<Spectrum, SamplePRD *, float *, float, float>(
        data->mat_ids[mat] + prd->params->material_base_offset, prd, pdf, u0, u1,
        MaterialCallMode::SAMPLE_FR);
    *pdf *= p;

    float true_pdf; // this is discarded
    return calcFr(prd, &true_pdf, data);
}

extern "C" __device__ Spectrum __direct_callable__mixed_material(SamplePRD *prd, float *pdf,
                                                                 float u0, float u1,
                                                                 MaterialCallMode mode)
{
    // get sbt data
    const uint mixed_instance_id = prd->inter.mat_instance_id;
    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    MixedMaterialData *data = reinterpret_cast<MixedMaterialData *>(sbt_record->callable_payload);
    data = data + mixed_instance_id;

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("%s\n", __FUNCTION__);
        printf("mesh id: %d\n", prd->inter.mesh_id);
        printf("mesh instance id: %d\n", prd->inter.mesh_instance_id);
        printf("mat id: %d\n", prd->inter.mat_id);
        printf("mat instance id: %d\n", prd->inter.mat_instance_id);
        printf("mat base offset: %d\n", prd->params->material_base_offset);
        printf("------------ sbt data ------------\n");
        printf("num_mat_mixed: %d\n", data->num_mat_mixed);
        printf("scale: %f\n", data->scale);
        printf("weights: ");
        for (size_t i = 0; i < data->num_mat_mixed; i++)
            printf("%f ", data->weights[i]);
        printf("\n");
        printf("mat_ids: ");
        for (size_t i = 0; i < data->num_mat_mixed; i++)
            printf("%d ", data->mat_ids[i]);
        printf("\n");
    }
#endif

    Spectrum result(0.0f);
    if (mode == MaterialCallMode::SAMPLE_FR)
    {
        result = rayGeneration(prd, pdf, data, u0, u1);
    }
    else if (mode == MaterialCallMode::CALC_FR_PDF)
    {
        result = calcFr(prd, pdf, data);
    }

    prd->inter.mat_instance_id = mixed_instance_id;

    return result;
}